#include <iostream>
#include <string>
#include <semaphore.h>
#include <unistd.h>
#include <sstream>
#include <vector>
#include <thread>
#include <mutex>
#include <string.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <sys/time.h>
#include <utility.h>

#ifdef __linux__
#include <pty.h>
#include <sys/epoll.h>
#elif defined(__APPLE__) || defined(__MACH__)
#include <util.h>
#include <sys/event.h>
#else
#error "Unsupported platform"
#endif
#include "mqtt_config.h"
#include "Belinda.h"
#define _set_error_code         \
    {                           \
        _error_code = __LINE__; \
        mux.unlock();           \
        return -__LINE__;       \
    }
using namespace std;
#define var auto
bool _client_connected = false;
int _client_sock_fd = 0;
mutex mux;
int _error_code = 0;
bool _mqtt_client_connected = false;
vector<string> _cmd_queue;
vector<string> _ext_queue;
int _got_terminal_info = 0; // used to send NAWS will
string _given_token;
void (*_on_socket_data)(char *, int) = NULL;
int _on_auth_stock_len = 0;
// return 0 if need to reset this handle
int (*_on_timeout_handle)() = NULL;
bool got_ttl = false;
int _max_request_count = 10;

void _telnet_send(unsigned char b1, unsigned char b2);
void _request_for_token();

#pragma region telnet
/*
telnet commands:
      SE                  240    End of subnegotiation parameters.
      NOP                 241    No operation.
      Data Mark           242    The data stream portion of a Synch. This should always be accompanied by a TCP Urgent notification.
      Break               243    NVT character BRK.
      Interrupt Process   244    The function IP.
      Abort output        245    The function AO.
      Are You There       246    The function AYT.
      Erase character     247    The function EC.
      Erase Line          248    The function EL.
      Go ahead            249    The GA signal.
      SB                  250    Indicates that what follows is subnegotiation of the indicated option.
      WILL (option code)  251    Indicates the desire to begin performing, or confirmation that you are now performing, the indicated option.
      WON'T (option code) 252    Indicates the refusal to perform, or continue performing, the indicated option.
      DO (option code)    253    Indicates the request that the other party perform, or confirmation that you are expecting the other party to perform, the indicated option.
      DON'T (option code) 254    Indicates the demand that the other party stop performing, or confirmation that you are no longer expecting the other party to perform, the indicated option.
      IAC                 255    Data Byte 255.

*/
const unsigned char IAC = 255;                   // \xff
const unsigned char SB = 250;                    // \xfa
const unsigned char SE = 240;                    // \xf0
const unsigned char NAWS = 31;                   // \x1f
const unsigned char DO = 253;                    // \xfd
const unsigned char DONT = 254;                  // \xfe
const unsigned char WILL = 251;                  // \xfb
const unsigned char WONT = 252;                  // \xfc
const unsigned char _ECHO = 1;                   // \x01
const unsigned char Suppress_Local_Echo = 3;     // \x03
const unsigned char GA = 249;                    // \xf9
const unsigned char Suppress_Go_Ahead = 3;       // \x03
const unsigned char LINEMODE = 34;               // \x22
const unsigned char Authentication = 37;         // \x25
const unsigned char Terminal_Type = 24;          // \x18
const unsigned char Terminal_Speed = 32;         // \x20
const unsigned char Remote_Flow_Control = 33;    // \x21
const unsigned char New_Environment_Option = 39; // \x27
const unsigned char Status = 5;                  // \x05
const unsigned char NewLine = 10;                // \x0a
const unsigned char CarriageReturn = 13;         // \x0d

void _telnet_send(unsigned char b1, unsigned char b2)
{
    unsigned char buf[3] = {IAC, b1, b2};
    send(_client_sock_fd, buf, 3, 0);
}
void _send_ext(string &ext)
{
    if (_mqtt_client_connected)
    {
        client_publish("ext", ext);
    }
    else
    {
        _ext_queue.push_back(ext);
    }
}
int go_WILL(unsigned char *b, int len)
{
    string msg;
    switch (b[0])
    {
    case NAWS:
        _telnet_send(DO, NAWS);
        _got_terminal_info = 1;
        return 2;
    case _ECHO:
        msg = "echo";
        _send_ext(msg);
        _telnet_send(DO, _ECHO);
        return 2;
    case Suppress_Go_Ahead:
        _telnet_send(DONT, Suppress_Go_Ahead);
        return 2;
    default:
        return 2;
    }
}
int go_WONT(unsigned char *b, int len)
{
    string msg;
    switch (b[0])
    {
    case NAWS:
        _telnet_send(DONT, NAWS);
        return 2;
    case _ECHO:
        msg = "-echo";
        _send_ext(msg);
        _telnet_send(DONT, _ECHO);
        return 2;
    case Suppress_Go_Ahead:
        _telnet_send(DONT, Suppress_Go_Ahead);
        return 2;
    default:
        return 2;
    }
}
int go_DO(unsigned char *b, int len)
{
    string msg;
    switch (b[0])
    {
    case NAWS:
        _telnet_send(WILL, NAWS);
        return 2;
    case _ECHO:
        msg = "echo";
        _send_ext(msg);
        // _telnet_send(WILL, _ECHO);
        return 2;
    case Suppress_Go_Ahead:
        // _telnet_send(WONT, Suppress_Go_Ahead);
        return 2;
    default:
        return 2;
    }
}
int go_DONT(unsigned char *b, int len)
{
    string msg;
    switch (b[0])
    {
    case NAWS:
        _telnet_send(WONT, NAWS);
        return 2;
    case _ECHO:
        msg = "-echo";
        _send_ext(msg);
        _telnet_send(WONT, _ECHO);
        return 2;
    case Suppress_Go_Ahead:
        // _telnet_send(WONT, Suppress_Go_Ahead);
        return 2;
    default:
        return 2;
    }
}
int go_SB(unsigned char *b, int len)
{
    if (b[0] == NAWS)
    {
        // NAWS
        if (len < 7)
            return 1;
        short cols = b[1] * 256 + b[2];
        short rows = b[3] * 256 + b[4];
        string msg = "size " + to_string(rows) + " " + to_string(cols);
        _send_ext(msg);
        _got_terminal_info = 2;
        return 8;
    }
    else
    {
        // skip the subnegotiation
        int i = 1;
        for (; i < len; i++)
        {
            if (b[i] == IAC && b[i + 1] == SE)
                break;
        }
        if (i >= len)
            return 1;
        return i + 2;
    }
}

#if DEBUG >= 2
#include <unordered_map>
void _telnet_translate_from_bytes(unsigned char *b, int n)
{
    static unordered_map<unsigned char, string> _map = {
        {IAC, ", IAC"},
        {WILL, "WILL"},
        {WONT, "WONT"},
        {DO, "DO"},
        {DONT, "DONT"},
        {SB, "SB"},
        {SE, "SE"},
        {GA, "GA"},
        {Suppress_Go_Ahead, "Suppress_Go_Ahead"},
        {LINEMODE, "LINEMODE"},
        {Authentication, "Authentication"},
        {Terminal_Type, "Terminal_Type"},
        {Terminal_Speed, "Terminal_Speed"},
        {Remote_Flow_Control, "Remote_Flow_Control"},
        {New_Environment_Option, "New_Environment_Option"},
        {Status, "Status"},
        {NAWS, "NAWS"},
        {_ECHO, "ECHO"},
        {NewLine, "NewLine"},
        {CarriageReturn, "CarriageReturn"},
    };
    for (int i = 0; i < n; i++)
    {
        if (_map.find(b[i]) != _map.end())
        {
            printf("%s ", _map[b[i]].c_str());
        }
        else
        {
            printf("%02x ", b[i]);
        }
    }
    printf("\r\n");
}
#endif

void _filter_data_from_sock(char *data, int len)
{
    static unsigned char _buf[1024 * 3];
    int i = 0, j = 0;
    unsigned char *b = (unsigned char *)data;

#if DEBUG >= 2
    if (*b == IAC)
    {
        printf("\x1b[0;32m");
        _telnet_translate_from_bytes(b, len);
        printf("\x1b[0m");
    }
#endif
    for (i = 0; i < len; i++)
    {
        if (b[i] != IAC)
        {
            _buf[j++] = b[i];
            continue;
        }
        i++;
        if (i >= len)
            break;
        if (b[i] == IAC)
        {
            _buf[j++] = b[i];
            continue;
        }
        if (b[i] == WILL)
        {
            i += go_WILL(b + i + 1, len - i - 1);
            i--;
            continue;
        }
        if (b[i] == WONT)
        {
            i += go_WONT(b + i + 1, len - i - 1);
            i--;
            continue;
        }
        if (b[i] == DO)
        {
            i += go_DO(b + i + 1, len - i - 1);
            i--;
            continue;
        }
        if (b[i] == DONT)
        {
            i += go_DONT(b + i + 1, len - i - 1);
            i--;
            continue;
        }
        if (b[i] == SB)
        {
            i += go_SB(b + i + 1, len - i - 1);
            i--;
            continue;
        }
    }

    if (j)
    {
        if (_mqtt_client_connected)
            mqtt_out_raw((char *)_buf, j);
        else
        {
            // push to _cmd_queue, and public when mqtt connected
            _cmd_queue.push_back(string((char *)_buf, j));
        }
    }

    if (!_got_terminal_info)
    {
        // send IAC DO NAWS, to get the terminal size
        _telnet_send(DO, NAWS);
        _got_terminal_info = 1;
    }
}
#pragma endregion

void _telnet_disconnected()
{
    _error_code = 0;
    _on_timeout_handle = NULL;
    mux.unlock();
}
// read from the socket, and send the data to the mqtt broker
#if __linux__
int _read_from_socket(int fd)
{
    int epoll_fd = epoll_create1(0);
    if (epoll_fd == -1)
    {
        _set_error_code;
    }
    struct epoll_event event;
    event.events = EPOLLIN | EPOLLRDHUP | EPOLLHUP | EPOLLET | EPOLLERR;
    event.data.fd = fd;

    if (epoll_ctl(epoll_fd, EPOLL_CTL_ADD, fd, &event) == -1)
    {
        _set_error_code;
    }
    const int MAX_EVENTS = 8;
    struct epoll_event events[MAX_EVENTS];
    const int buf_len = 1024 * 3;
    static char buf[buf_len];
    while (1)
    {
        int num_events = epoll_wait(epoll_fd, events, MAX_EVENTS, -1);
        if (num_events == -1)
        {
            cerr << "epoll_wait error" << endl;
            _set_error_code;
        }
        for (int i = 0; i < num_events; i++)
        {
            if (events[i].events & EPOLLIN)
            {

                int num_read = read(events[i].data.fd, buf, buf_len);
                if (num_read == -1)
                {
                    close(epoll_fd);
                    _set_error_code;
                }
                if (num_read == 0)
                {
                    if (epoll_ctl(epoll_fd, EPOLL_CTL_DEL, events[i].data.fd, NULL) == -1)
                    {
                        close(epoll_fd);
                        _set_error_code;
                    }
                    close(epoll_fd);
                    _telnet_disconnected();
                    return 0;
                }
                else
                {
                    if (_client_connected && _on_socket_data)
                        _on_socket_data(buf, num_read);
                }
            }
            if (events[i].events & (EPOLLRDHUP | EPOLLHUP | EPOLLERR))
            {
                if (epoll_ctl(epoll_fd, EPOLL_CTL_DEL, events[i].data.fd, NULL) == -1)
                {
                    close(epoll_fd);
                    _set_error_code;
                }
                close(epoll_fd);
                _telnet_disconnected();
                return 0;
            }
        }
    }
    // unknown reson
    close(epoll_fd);
    _error_code = 1;
    mux.unlock();
    return 0;
}
#else
int _read_from_socket(int fd)
{
    int kq = kqueue();
    if (kq == -1)
    {
        _set_error_code;
    }
    struct kevent event;
    EV_SET(&event, fd, EVFILT_READ, EV_ADD, 0, 0, NULL);

    if (kevent(kq, &event, 1, NULL, 0, NULL) == -1)
    {
        _set_error_code;
    }
    const int MAX_EVENTS = 8;
    struct kevent events[MAX_EVENTS];
    const int buf_len = 1024 * 3;
    static char buf[buf_len];
    while (1)
    {
        int num_events = kevent(kq, NULL, 0, events, MAX_EVENTS, NULL);
        if (num_events == -1)
        {
            break;
        }
        for (int i = 0; i < num_events; i++)
        {
            if (events[i].filter == EVFILT_READ)
            {
                int num_read = read(events[i].ident, buf, buf_len);
                if (num_read == -1)
                {
                    close(kq);
                    _set_error_code;
                }
                if (num_read == 0)
                {
                    close(kq);
                    _telnet_disconnected();
                    return 0;
                }
                else
                {
                    if (_client_connected && _on_socket_data)
                        _on_socket_data(buf, num_read);
                }
            }
        }
    }
    close(kq);
    _error_code = 1;
    mux.unlock();
    return 0;
}
#endif

time_t get_current_time_ms()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    return tv.tv_sec * 1000 + tv.tv_usec / 1000;
}
void _submit_password(string &password)
{
    if (_given_token.empty())
        return;
    auto auth = _given_token + password;
    auto msg = get_md5_str(auth.c_str(), auth.length());
    msg = "auth:" + msg;
    client_publish("ext", msg);
}
void _on_auth(char *data, int len)
{
    static char _stock[32] = "";
    bool type_ok = false;
    if (len == 0)
        return;
    unsigned char *b = (unsigned char *)data;
    for (int i = 0; i < len; i++)
    {
        if (b[0] == IAC) // like: IAC DO ECHO
        {
            i += 2;
            continue;
        }
        if (_on_auth_stock_len >= 1 && _stock[_on_auth_stock_len - 1] == '\r' && data[i] == '\n')
        {
            _stock[_on_auth_stock_len - 1] = '\0';
            _on_auth_stock_len--;
            type_ok = true;
            break;
        }
        _stock[_on_auth_stock_len++] = data[i];
        if (_on_auth_stock_len >= 32)
        {
            type_ok = true;
            break;
        }
    }
    if (!type_ok)
        return;
    string msg;
    if (_given_token.empty())
    {
        msg = "reqst";
        int rc = client_publish("ext", msg);
        send_by_socket_s("\r\n", 0);

        return;
    }
    string auth(_stock, _on_auth_stock_len);
    send_by_socket_s("\r\n", 0);
    _submit_password(auth);
}

void on_message(char *topic, int topic_len, char *msg, int msg_len)
{
#if DEBUG >= 2
    cout << "MSG [" << string(topic, topic_len) << "]:" << endl;
    // print the msg in bytes
    for (int i = 0; i < msg_len; i++)
    {
        char c = msg[i];
        if (c == '\r')
            cout << "\\r";
        else if (c == '\n')
            cout << "\\n";
        else if (c == '\t')
            cout << "\\t";
        else if (c == '\0')
            cout << "\\0";
        else if (c >= 32 && c <= 126)
            cout << c;
        else
            cout << "\\x" << hex << (int)c;
    }
    cout << endl;
#endif
    if (topic_len > 4 && strncmp(topic + topic_len - 4, "/msg", 4) == 0)
    {
        if (_client_connected)
            send(_client_sock_fd, msg, msg_len, 0);
        return;
    }

    if (topic_len > 5 && strncmp(topic + topic_len - 5, "/pong", 5) == 0)
    {
        if (msg_len == 0)
            return;
        if (*msg == '#')
        {
            _given_token = string(msg + 1, msg_len - 1);
            auto password = get_client_password();
            if (!password.empty())
            {

                _submit_password(password);
                return;
            }
            send_by_socket_s("Enter password>", 'g');
            _telnet_send(WILL, _ECHO);
            _on_auth_stock_len = 0;
            return;
        }
        if (*msg == '!') // failed to auth, exit the client
        {
            send_by_socket_s("Authentication failed, halted!\r\n", 'r');
            _error_code = 0;
            mux.unlock();
            return;
        }
        if (*msg == '>') // failed to auth, try again
        {
            _given_token = string(msg + 1, msg_len - 1);
            // if use password in config file, this means it's incorrect and then ask user to input
            send_by_socket_s("Enter password>", 'g');
            _telnet_send(WILL, _ECHO);
            _on_auth_stock_len = 0;
            return;
        }
        if (*msg == '$') // pass the auth
        {
            send_by_socket_s("[Login]\r\n", 'g');
            _telnet_send(WONT, _ECHO);
            _on_socket_data = _filter_data_from_sock;
            for (auto &cmd : _cmd_queue)
            {
                mqtt_out_raw(cmd);
            }
            _cmd_queue.clear();
            for (auto &ext : _ext_queue)
            {
                client_publish("ext", ext);
            }
            _ext_queue.clear();
            return;
        }
        // show TTL there
        got_ttl = true;
        auto ms = get_current_time_ms();
        auto ttl = ms - stol(string(msg, msg_len));
        string ttl_msg = " TTL: " + to_string(ttl) + " ms \r\n";
        send_by_socket(ttl_msg, 'g');
        if (_given_token.empty())
        {
            _request_for_token();
        }
        return;
    }
}

void _measure_TTL()
{
    string time;
    // ping the current time
    auto ms = get_current_time_ms();
    time = to_string(ms);
    client_publish("ping", time);
}
void _request_for_token()
{
    string request = "reqst";
    client_publish("ext", request); // request the token
}
void __init_each_connection()
{
    _got_terminal_info = 0;
    _on_socket_data = _on_auth;
    _max_request_count = 10;
    _on_timeout_handle = NULL;
    got_ttl = false;
    _given_token = "";
}
int _request_for_token_and_ttl()
{
    _max_request_count--;
    if (_max_request_count <= 0)
    {
        return 0;
    }
    bool clear = true;
    if (!got_ttl)
    {
        _measure_TTL();
        clear = false;
        usleep(1000 * (500 + _max_request_count * 100));
    }
    if (got_ttl && _given_token.empty())
    {
        // usleep(1000 * (500 + _max_request_count * 50));
        string request = "reqst";
        client_publish("ext", request); // request the token
        clear = false;
        usleep(1000 * (500 + _max_request_count * 50));
    }
    if (clear)
        return 0;
    mux.unlock();
    return 1;
}
int _request_for_token_after_connect()
{
    _measure_TTL();
    // _request_for_token();

    // send the command queue, if any
    for (auto &cmd : _cmd_queue)
    {
        mqtt_out_raw(cmd);
    }
    for (auto &ext : _ext_queue)
    {
        client_publish("ext", ext);
    }
    _cmd_queue.clear();
    _ext_queue.clear();
    _on_timeout_handle = _request_for_token_and_ttl;
    mux.unlock();
    return 1;
}

void on_error(int error_code)
{

    if (error_code)
    {
        // rasie an exception, and exit the master process
        _error_code = 1;
        _mqtt_client_connected = false;
        mux.unlock();
        return;
    }
    _mqtt_client_connected = true;
    send_by_socket_s("MQTT Broker Connected\r\n", 'g');
    string request = "reqst";
    int rc = client_publish("ext", request); // request the token
    if (_on_socket_data != _filter_data_from_sock)
    {
        if (_client_connected)
        {
            _on_socket_data = _on_auth;
        }
    }
    _on_timeout_handle = _request_for_token_after_connect;
    mux.unlock();
    // start in master process
}

void send_by_socket(string &msg, char color)
{
    send_by_socket(msg.c_str(), msg.length(), color);
}

void send_by_socket_s(const char *msg, char color)
{
    if (!msg)
        return;
    int len = strlen(msg);
    if (len == 0)
        return;
    send_by_socket(msg, len, color);
}

void send_by_socket(const char *msg, int len, char color)
{
    if (!_client_connected || !msg || !len)
    {
        return;
    }
    if (color == 0)
    {
        send(_client_sock_fd, msg, len, 0);
        return;
    }
    const char *red = "\033[31m";
    const char *green = "\033[32m";
    const char *yellow = "\033[33m";
    const char *blue = "\033[34m";
    const char *magenta = "\033[35m";
    const char *cyan = "\033[36m";
    const char *white = "\033[37m";
    const char *reset = "\033[0m";
    const char *use_color = NULL;
    switch (color)
    {
    case 'r':
        use_color = red;
        break;
    case 'g':
        use_color = green;
        break;
    case 'y':
        use_color = yellow;
        break;
    case 'b':
        use_color = blue;
        break;
    case 'm':
        use_color = magenta;
        break;
    case 'c':
        use_color = cyan;
        break;
    case 'w':
        use_color = white;
        break;
    default:
        use_color = NULL;
        break;
    }
    if (use_color)
    {
        send(_client_sock_fd, use_color, strlen(use_color), 0);
    }
    send(_client_sock_fd, msg, len, 0);
    if (use_color)
    {
        send(_client_sock_fd, reset, strlen(reset), 0);
    }
}

/*
  usage: ./Belinda [port] [config_file]
*/
const string _main_usage = "usage: Belinda [-h] [-v] [port] [config_file]";
int main(int argn, char *argv[])
{
    if (argn > 1 && (strcmp(argv[1], "-h") == 0 || strcmp(argv[1], "--help") == 0))
    {
        cout << _main_usage << endl;
        return 0;
    }
    if (argn > 1 && (strcmp(argv[1], "-v") == 0 || strcmp(argv[1], "--version") == 0))
    {
        cout << "Belinda version: " << Woojie_VERSION << "\nbuild time:" << BUILD_TIMESTAMP << endl;
        return 0;
    }
    int server_fd, valread;
    struct sockaddr_in address;
    int opt = 1;
    int addrlen = sizeof(address);
    char buffer[1024] = {0};
    int listen_port = 23;
    int _n = 1;
    if (argn > _n && isdigit(argv[1][0]))
    {
        listen_port = atoi(argv[1]);
        _n++;
    }
    string confg_file = "mqtt.ini";
    if (argn > _n)
        confg_file = argv[_n];
    int rc = 0;
    rc = read_config(confg_file);
    if (rc)
        exit(rc);
    server_fd = socket(AF_INET, SOCK_STREAM, 0);
    if (server_fd <= 0)
    {
        perror("socket failed");
        exit(__LINE__);
    }
#if __linux__
    // Forcefully attaching socket to the port
    if (setsockopt(server_fd, SOL_SOCKET, SO_REUSEADDR | SO_REUSEPORT, &opt, sizeof(opt)) == -1)
    {
        perror("setsockopt");
        exit(__LINE__);
    }
#endif
    address.sin_family = AF_INET;
    address.sin_addr.s_addr = INADDR_ANY;
    address.sin_port = htons(listen_port);

    if (::bind(server_fd, (struct sockaddr *)&address,
               sizeof(address)) < 0)
    {
        perror("bind failed");
        exit(__LINE__);
    }
    if (listen(server_fd, 1) < 0)
    {
        perror("listen");
        exit(__LINE__);
    }
    thread t1;
    set_on_message_callback(on_message);
    set_on_error_callback(on_error);
    try
    {
        while (1)
        {
            if ((_client_sock_fd = accept(server_fd, (struct sockaddr *)&address,
                                          (socklen_t *)&addrlen)) < 0)
            {
                perror("accept");
                close(server_fd);
                exit(__LINE__);
            }
#ifdef DEBUG
            cout << "client " << inet_ntoa(address.sin_addr) << ":" << ntohs(address.sin_port) << " connected" << endl;
#endif
            _client_connected = true;
            __init_each_connection();
            mux.lock();
            start_loop();
// create a thread to listen to the socket
#if __linux__
            t1 = thread(_read_from_socket, _client_sock_fd);
#else
            t1 = thread([]()
                        { _read_from_socket(_client_sock_fd); });
#endif
            t1.detach();
            // wait the mutex to be released
            while (1)
            {
                mux.lock();
                if (_error_code)
                {
                    throw std::runtime_error("error in the thread:" + to_string(_error_code));
                }
                if (!_on_timeout_handle)
                    break;
                int rc = _on_timeout_handle();
                if (!rc)
                    _on_timeout_handle = NULL;
            }

            // after the aquire the mux, close the socket and mqtt client
            _client_connected = false;
            _given_token = "";
            close(_client_sock_fd);
            string msg = "exit";
            int rc = client_publish("ext", msg); // notify to close pty
            if (rc)
                cout << "notify exit:" << rc << endl;
            stop_loop();
            _mqtt_client_connected = false;
            if (t1.joinable())
                t1.join();
            mux.unlock();
#ifdef DEBUG
            cout << "client disconnected" << endl;
#endif
        }
    }
    catch (const std::exception &e)
    {
        std::cerr << e.what() << '\n';
    }
    if (_client_connected)
    {
        _client_connected = false;
        close(_client_sock_fd);
        stop_loop();
        if (t1.joinable())
            t1.join();
    }
    close(server_fd);
    mux.try_lock();
    mux.unlock();
    return 0;
}